package com.timeriver.cases.power_prediction.v4

import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.ml.regression.{RandomForestRegressionModel, RandomForestRegressor}
import org.apache.spark.sql.{DataFrame, SparkSession}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.{DateType, DoubleType, IntegerType}

/**
  * 使用集成模型（随机森林）对特征重要性进行排序:
  *   dayOfWeek:0.1357248256579423
  *   weekOfYear:0.2922662695591027
  *   weekdayOrNot:0.022197062415602015
  *   dayOfMonth:0.19858240974192642
  *   tenDays:0.043023803153029214
  * 由此可知：weekdayOrNot、tenDays特征可以删除
  */
object PowerRegressionPredict {
  def main(args: Array[String]): Unit = {
    val session: SparkSession = SparkSession.builder()
      .appName("基于随机森林回归模型进行特征重要度分析")
      .master("local[12]")
      .getOrCreate()

    val rawData: DataFrame = session.read.format("csv")
      .option("header", true)
      .load("D:\\workspace\\gitee_space\\spark-ml-machine-learning\\data\\zhenjiang_power.csv")

    val data: DataFrame = rawData.withColumn("record_date", col("record_date").cast(DateType))
      .withColumn("power_consumption", col("power_consumption").cast(DoubleType))
      .withColumn("user_id", col("user_id").cast(IntegerType))

    /** 地区每日耗电量 */
    val daily_power: DataFrame = data.groupBy("record_date").sum("power_consumption").sort("record_date")

    /** 构造数据特征：月份/星期几/一年的第几周/是否是周末/一月中的上中下旬 */
    val newData: DataFrame = daily_power.withColumn("month", month(column("record_date")))
      .withColumn("dayOfWeek", dayofweek(column("record_date")))
      .withColumn("weekOfYear", weekofyear(column("record_date")))
      .withColumn("dayOfMonth", dayofmonth(column("record_date")))
      .withColumnRenamed("sum(power_consumption)", "label")
      .selectExpr("*",
        "case when dayOfWeek=6 then 1 when dayOfWeek=7 then 1 else 0 end as weekdayOrNot",
        "case when dayOfMonth between 1 and 10 then 1 when dayOfMonth between 11 and 20 then 2 else 3 end as tenDays"
      )

    val trainData: DataFrame = new VectorAssembler()
      .setInputCols(Array("month", "dayOfWeek", "weekOfYear", "weekdayOrNot", "dayOfMonth", "tenDays"))
      .setOutputCol("features")
      .transform(newData)

    val regressor: RandomForestRegressor = new RandomForestRegressor()
      .setFeaturesCol("features")
      .setLabelCol("label")
      .setMaxBins(32)
      .setMaxDepth(15)

    val model: RandomForestRegressionModel = regressor.fit(trainData)

    println(s"特征重要度分析：${model.featureImportances}")
  }
}
